import torch
from torch import nn

from torch.autograd import Variable

from fpnNet import FPN50, FPN101
from config import config

class retinaNet(nn.Module):
    _num_anchors    = config.network.num_anchors

    def __init__(self, num_classes=config.dataset.num_classes):
        """
        initialization of the detection net
        """
        super(retinaNet, self).__init__()
        self._num_classes   = num_classes 
        self._init_net()

    def _init_net(self):
        if config.network.backbone      == "FPN50":
            self._fpn   = FPN50()
        elif config.network.backbone    == "FPN101":
            self._fpn   = FPN101()
        self._loc_head  = self._make_head(self._num_anchors * 4)
        self._cls_head  = self._make_head(self._num_anchors * self._num_classes)

    def _make_head(self, out_planes):
        """
        make head layer for location or classification
        """
        layers  = []
        for i in range(4):
            _i_conv = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
            _i_relu = nn.ReLU(inplace=True)
            layers.append(_i_conv)
            layers.append(_i_relu)
        _last_conv  = nn.Conv2d(256, out_planes, kernel_size=3, stride=1, padding=1)
        layers.append(_last_conv)
        return nn.Sequential(*layers)

    def forward(self, x):
        """
        forward to do detection
        """
        batch_size  = x.size()[0]
        # firstly extract conv features
        fms     = self._fpn(x)
        _loc_pred_list  = []
        _cls_pred_list  = []
        for i in range(len(fms)):
            # calculate classification and location prediction using each conv feature map
            _i_fm       = fms[i]
            _i_loc_pred = self._loc_head(_i_fm)
            _i_cls_pred = self._cls_head(_i_fm)
            _i_loc_pred = _i_loc_pred.permute(0,2,3,1).contiguous().view(batch_size,-1,4)
            _i_cls_pred = _i_cls_pred.permute(02,3,1).contiguous().view(batch_size,-1,self._num_classes)
            _loc_pred_list.append(_i_loc_pred)
            _cls_pred_list.append(_i_cls_pred)
        _loc_pred_var   = torch.cat(_loc_pred_list, 1)
        _cls_pred_var   = torch.cat(_cls_pred_list, 1)
        return _loc_pred_var, _cls_pred_var

def test():
    net     = RetinaNet()
    inp     = Variable(torch.randn(2,3,224,224))
    if config.useGpu:
        net = net.cuda()
        inp = inp.cuda()
    loc_pred_var, cls_pred_var  = net(inp)
    
    # test backward
    loc_pred_grad, cls_pred_grad= Variable(torch.randn(loc_pred_var.size())), Variable(torch.randn(cls_pred_var.size()))
    if config.useGpu:
        loc_pred_grad, cls_pred_grad    = loc_pred_grad.cuda(), cls_pred_grad.cuda()
    loc_pred_var.backward(loc_pred_grad)
    cls_pred_var.backward(cls_pred_grad)
